import os

import numpy as np
import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置中文字体为黑体
plt.rcParams['axes.unicode_minus'] = False
def calculate_utility(user_pair_idx, hap_idx, connection, h_hap2user, user_pairs, P_hap, N_max, noise, B):
    connected_pairs = np.sum(connection[hap_idx, :])
    if connected_pairs == 0:
        return 0.0  # 如果HAP没有连接任何用户对，返回0效用

    M = connection.shape[0]  # HAP的数量
    B_pair = B / M / connected_pairs
    P_pair = P_hap / connected_pairs
    # 获取用户对中的两个用户
    user_p, user_q = user_pairs[user_pair_idx]

    # 获取信道系数
    h_p = h_hap2user[hap_idx, int(user_p)]
    h_q = h_hap2user[hap_idx, int(user_q)]
    alpha_p = 0.33  # 功率分配系数
    alpha_q = 0.67

    # 计算SINR和速率
    sinr_p = (alpha_p * P_pair * h_p ** 2) / noise
    sinr_q_p = (alpha_q * P_pair * h_q ** 2) / (alpha_p * P_pair * h_q ** 2 + noise)

    # 计算用户速率
    rate_p = B_pair * np.log2(1 + sinr_p)
    rate_q = B_pair * np.log2(1 + sinr_q_p)

    # 总效用为两个用户的速率之和
    utility = rate_p + rate_q

    return utility
def calculate_system_rate(connection, h_hap2user, user_pairs, P_hap, N_max, noise, B, use_optimal_power):
    M, N_total = connection.shape
    N = user_pairs.shape[0]
    sum_rate = 0.0
    for hap_idx in range(M):  # 遍历每个HAP
        # 获取当前HAP连接的真实用户对
        connected_pairs = np.where(connection[hap_idx, :N] == 1)[0]
        num_connected = len(connected_pairs)

        if num_connected == 0:
            continue  # 如果没有连接任何用户对，跳过

        # 分配给每个用户对的带宽和功率
        B_pair = B / M / num_connected
        P_pair = P_hap / num_connected

        # 遍历每个连接的用户对
        for j in range(num_connected):
            pair_idx = connected_pairs[j]

            # 获取两个用户索引
            user_p, user_q = user_pairs[pair_idx].astype(int)

            # 获取信道增益
            h_p = h_hap2user[hap_idx, user_p]
            h_q = h_hap2user[hap_idx, user_q]

            # 根据是否使用优化功率分配
            if use_optimal_power:
                # 优化功率分配：根据信道强度交换用户并计算 alpha
                if h_p < h_q:
                    h_p, h_q = h_q, h_p
                    user_p, user_q = user_q, user_p

                beta = h_q / h_p

                if beta < 0.5:
                    alpha_p = (1 - beta) / (1 + beta)
                    alpha_q = 2 * beta / (1 + beta)
                else:
                    alpha_p = 1 / (1 + 2 * beta)
                    alpha_q = 2 * beta / (1 + beta)

                total_alpha = alpha_p + alpha_q
                alpha_p /= total_alpha
                alpha_q /= total_alpha

            else:
                # 固定功率分配
                alpha_p = 0.33
                alpha_q = 0.67

            # 计算SINR
            sinr_p = (alpha_p * P_pair * h_p ** 2) / noise
            sinr_q_p = (alpha_q * P_pair * h_q ** 2) / (alpha_p * P_pair * h_q ** 2 + noise)

            # 计算速率
            rate_p = B_pair * np.log2(1 + sinr_p)
            rate_q = B_pair * np.log2(1 + sinr_q_p)

            # 累加总速率
            sum_rate += rate_p + rate_q

    return sum_rate
def NOMA(M,P_hap,output_dir):
    # 定义参数
    M = M                    # HAP数量
    B = 100e6                # 总带宽（Hz）
    P_s = 40                 # 卫星向HAP发送信号的功率（dBm）
    P_s_linear = 10 ** (P_s / 10 - 3)  # 转换为线性值（W）

    N_max = 15               # 每个HAP可服务的NOMA用户对数
    N = 50                   # 实NOMA用户对总数
    N_virtual = 25           # 虚拟NOMA用户对数

    total_pairs = M * N_max  # 总用户对数量为75

    P_hap = P_hap               # 一个HAP向地面用户发送信号的总功率（dBm）
    P_hap_linear = 10 ** (P_hap / 10 - 3)  # 转换为线性值（W）

    noise_power = -130       # 噪声功率（dBm）
    noise_linear = 10 ** (noise_power / 10 - 3)  # 噪声功率线性值（W）
    hap_altitude = 20000     # HAP高度（m）
    ground_area = 1000       # 地面区域大小 (1000 x 1000)

# =================== 生成HAP位置 ===================
    hap_positions = np.zeros((M, 3))

    for i in range(M):
        angle = 2 * np.pi / M * (i + 1)  # MATLAB 中索引从1开始，所以是 i+1
        x = 500 * np.cos(angle) + 500
        y = 500 * np.sin(angle) + 500
        z = hap_altitude
        hap_positions[i, :] = [x, y, z]


# =================== 生成地面用户位置 ===================
    user_positions = np.zeros((2 * N, 3))

    for i in range(2 * N):
        user_positions[i, :] = [
            ground_area * np.random.rand(),
            ground_area * np.random.rand(),
            0  # 地面高度为0
        ]
    # =================== 生成卫星到HAP的信道系数 ===================
    h_s2hap = 1.0 + 8.9 * np.random.rand(M, 1)
    # HAP到真实用户的信道系数
    h_hap2user = 0.1 + 0.8 * np.random.rand(M, 2*N)
    remaining_users = list(range(2 * N))  # 所有用户的索引 [0, 1, ..., 99]

    import matlab.engine


    # 启动 MATLAB 引擎
    eng = matlab.engine.start_matlab()

    # 添加 NOMA.m 所在的目录（可选）
    # eng.addpath(r'C:\your\matlab\script\path', nargout=0)

    # 调用 MATLAB 函数 NOMA
    user_pairs, user_positions = eng.NOMA(nargout=2)
    user_pairs_np = np.array(user_pairs, dtype=np.float64) -1 # 先转 float，再转 int
    user_pairs = np.int32(user_pairs_np)

    user_positions = np.array(user_positions)
    # 关闭引擎
    eng.quit()

    connection = np.zeros((M, total_pairs), dtype=int)
    # 记录每个用户对分配的HAP索引（从0开始）
    user_hap_assignment = np.zeros(total_pairs, dtype=int)
    # 记录每个HAP当前分配的用户对数量
    hap_load = np.zeros(M, dtype=int)
    for j in range(total_pairs):
        # 找出未满载的 HAP
        valid_haps = [i for i in range(M) if hap_load[i] < N_max]

        if not valid_haps:
            raise ValueError('无法满足所有HAP约束，请增加HAP数量或减少用户对数量')

        # 随机选择一个 HAP
        selected_hap = np.random.choice(valid_haps)

        # 更新记录
        user_hap_assignment[j] = selected_hap
        hap_load[selected_hap] += 1
        connection[selected_hap, j] = 1
    rand_connection = np.copy(connection)
    max_iterations = 100
    iter = 0
    exchange_occurred = True
    while exchange_occurred and iter < max_iterations:
        iter += 1
        exchange_occurred = False

        for j1 in range(N):
            i1 = np.where(rand_connection[:, j1] == 1)[0]  # 当前用户对j1的HAP

            if len(i1) == 0:
                continue  # 确保用户对分配了HAP

            i1 = i1[0]  # 取出具体的索引值

                # 计算当前用户对j1的效用
            current_utility_j1 = calculate_utility(j1, i1, rand_connection, h_hap2user, user_pairs, P_hap_linear, N_max,
                                                       noise_linear, B)

                # 尝试与其他用户对交换（包括虚拟用户对）
            for j2 in range(total_pairs):
                if j1 == j2:
                    continue  # 跳过自己

                i2 = np.where(rand_connection[:, j2] == 1)[0]  # 用户对j2的HAP

                if len(i2) == 0 or i1 == i2[0]:
                    continue  # 确保j2有HAP且与j1的HAP不同

                i2 = i2[0]  # 取出具体的索引值

                    # 计算当前用户对j2的效用
                if j2 < N:  # 真实用户对
                    current_utility_j2 = calculate_utility(j2, i2, rand_connection, h_hap2user, user_pairs, P_hap_linear,
                                                               N_max, noise_linear, B)
                else:  # 虚拟用户对
                    current_utility_j2 = 0  # 虚拟用户对的效用为0

                    # 交换连接并计算新效用
                new_connection = np.copy(connection)
                new_connection[i1, j1] = 0
                new_connection[i2, j1] = 1
                new_connection[i2, j2] = 0
                new_connection[i1, j2] = 1

                    # 计算交换后的效用
                new_utility_j1 = calculate_utility(j1, i2, new_connection, h_hap2user, user_pairs, P_hap_linear, N_max,
                                                       noise_linear, B)

                if j2 < N:  # 真实用户对
                    new_utility_j2 = calculate_utility(j2, i1, new_connection, h_hap2user, user_pairs, P_hap_linear, N_max,
                                                           noise_linear, B)
                else:  # 虚拟用户对
                    new_utility_j2 = 0  # 虚拟用户对的效用为0

                    # 判断是否为交换阻塞对
                if (new_utility_j1 > current_utility_j1) and (j2 >= N or new_utility_j2 >= current_utility_j2):
                        # 交换被批准
                    connection = np.copy(new_connection)
                    exchange_occurred = True
                    break  # 找到交换阻塞对后跳出内循环

            if exchange_occurred:
                break  # 找到交换阻塞对后跳出外循环，重新开始


        opt_connection = np.copy(connection)
    nearest_connection = np.zeros((M, total_pairs), dtype=int)
    hap_load = np.zeros(M, dtype=int)  # 记录每个HAP当前负载

    for j in range(N):
        user1_idx = int(user_pairs[j, 0])
        user2_idx = int(user_pairs[j, 1])

        # 获取用户的二维位置（忽略高度）
        user1_pos = np.array([user_positions[user1_idx, 0], user_positions[user1_idx, 1]])
        user2_pos = np.array([user_positions[user2_idx, 0], user_positions[user2_idx, 1]])

        avg_distances = []

        # 遍历每个 HAP，计算平均距离
        for i in range(M):
            hap_pos = hap_positions[i, :2]  # 只取 x, y 坐标
            dist1 = np.linalg.norm(hap_pos - user1_pos)
            dist2 = np.linalg.norm(hap_pos - user2_pos)
            avg_dist = (dist1 + dist2) / 2
            avg_distances.append((avg_dist, i))  # 存储距离和对应的HAP索引

        # 按距离从小到大排序 HAP
        avg_distances.sort(key=lambda x: x[0])
        sorted_haps = [i for (d, i) in avg_distances]

        # 寻找一个负载未满的 HAP
        for hap_idx in sorted_haps:
            if hap_load[hap_idx] < N_max:
                nearest_connection[hap_idx, j] = 1
                hap_load[hap_idx] += 1
                break
    for j in range(N, total_pairs):
        # 找出负载未满的 HAP
        valid_haps = [i for i in range(M) if hap_load[i] < N_max]

        if not valid_haps:
            continue  # 如果没有可用HAP，则跳过此用户对

        # 随机选择一个 HAP
        selected_hap = np.random.choice(valid_haps)

        # 更新连接矩阵和负载
        nearest_connection[selected_hap, j] = 1
        hap_load[selected_hap] += 1
    nearest_avg_sum_rate = calculate_system_rate(nearest_connection, h_hap2user, user_pairs, P_hap_linear, N_max, noise_linear, B, False)


    opt_avg_sum_rate = calculate_system_rate(opt_connection, h_hap2user, user_pairs, P_hap_linear, N_max, noise_linear, B, False)


    opt_opt_sum_rate = calculate_system_rate(opt_connection, h_hap2user, user_pairs, P_hap_linear, N_max, noise_linear, B, True)
    plt.figure(figsize=(10, 6))

# 绘制实用户位置（蓝色实心圆点）
    plt.scatter(user_positions[:2 * N, 0], user_positions[:2 * N, 1],
            s=30, c='b', marker='o', label='实用户', edgecolors='none')

# 绘制 HAP 投影位置（绿色实心菱形）
    plt.scatter(hap_positions[:, 0], hap_positions[:, 1],
            s=100, c='g', marker='d', label='HAP（投影）', edgecolors='none')

# 连接配对用户
    for i in range(N):
        user1_idx = user_pairs[i, 0]  # Python 索引从0开始
        user2_idx = user_pairs[i, 1]

        # 获取用户的二维位置（忽略高度）
        user1_pos = user_positions[user1_idx, :2]  # 直接取前两个元素 x, y
        user2_pos = user_positions[user2_idx, :2]

        # 绘制连接线
        plt.plot([user1_pos[0], user2_pos[0]],
                 [user1_pos[1], user2_pos[1]],
                 'k-', linewidth=0.5)

    # 添加标签和图
    plt.title('用户分布图')
    plt.xlabel('X坐标')
    plt.ylabel('Y坐标')
    plt.legend()
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.axis('equal')  # 保持坐标轴比例一致
    plt.tight_layout()
    # 完整的输出文件路径
    output_filename = os.path.join(output_dir, "plot_users.png")
    # 保存图像
    plt.savefig(output_filename, dpi=300, bbox_inches='tight')
    # 显示图像
    plt.show()
    snr_db = np.arange(-10, 31, 5)
    snr_linear = 10 ** (snr_db / 10)
    # 初始化结果数组
    rates_nearest_avg = np.zeros_like(snr_db, dtype=float)
    rates_opt_avg = np.zeros_like(snr_db, dtype=float)
    rates_opt_opt = np.zeros_like(snr_db, dtype=float)
    for i in range(len(snr_db)):
        P_current = snr_linear[i] * noise_linear

        rates_nearest_avg[i] = calculate_system_rate(
            nearest_connection, h_hap2user, user_pairs,
            P_current, N_max, noise_linear, B, use_optimal_power=False
        )
        rates_opt_avg[i] = calculate_system_rate(
            opt_connection, h_hap2user, user_pairs,
            P_current, N_max, noise_linear, B, use_optimal_power=False
        )
        rates_opt_opt[i] = calculate_system_rate(
            opt_connection, h_hap2user, user_pairs,
            P_current, N_max, noise_linear, B, use_optimal_power=True
        )

    plt.figure(figsize=(10, 6))

    # 绘制三条曲线
    plt.plot(snr_db, rates_nearest_avg, 'b-o', linewidth=2, label='就近匹配+固定功率分配')
    plt.plot(snr_db, rates_opt_avg, 'r-s', linewidth=2, label='优化匹配+固定功率分配')
    plt.plot(snr_db, rates_opt_opt, 'g-d', linewidth=2, label='优化匹配+优化功率分配')

    # 设置标题和坐标轴标签
    plt.title('系统总和速率随HAP处SNR变化', fontsize=14)
    plt.xlabel('SNR (dB)', fontsize=12)
    plt.ylabel('总和速率 (bps)', fontsize=12)

    # 显示图例
    plt.legend(fontsize=12)

    # 启用网格线
    plt.grid(True, linestyle='--', alpha=0.7)

    # 自动调整布局以防止重叠
    plt.tight_layout()


    # 完整的输出文件路径
    output_filename = os.path.join(output_dir, "plot_HAP.png")

    # 保存图像
    plt.savefig(output_filename, dpi=300, bbox_inches='tight')
# 显示图像
    plt.show()

    # 定义用户对数量范围
    user_counts = np.arange(20, 51, 5)  # 从20到50，步长为5
    # 初始化结果数组
    rates_nearest_avg_users = np.zeros(len(user_counts))
    rates_opt_avg_users = np.zeros(len(user_counts))
    rates_opt_opt_users = np.zeros(len(user_counts))
    # 固定 SNR = 20 dB
    fixed_snr_db = 20
    fixed_P = 10 ** (fixed_snr_db / 10) * noise_linear  # 使用之前定义的 noise_linear
    for i in range(len(user_counts)):
        current_user_count = user_counts[i]

        # 确保不超过最大用户对数量
        if current_user_count > user_pairs.shape[0]:
            current_user_count = user_pairs.shape[0]

        # 截取当前用户数量对应的连接矩阵和用户对
        current_nearest = nearest_connection[:, :current_user_count]
        current_opt = opt_connection[:, :current_user_count]
        current_pairs = user_pairs[:current_user_count, :]

        # 计算系统总速率
        rates_nearest_avg_users[i] = calculate_system_rate(
            current_nearest, h_hap2user, current_pairs,
            fixed_P, N_max, noise_linear, B, use_optimal_power=False
        )
        rates_opt_avg_users[i] = calculate_system_rate(
            current_opt, h_hap2user, current_pairs,
            fixed_P, N_max, noise_linear, B, use_optimal_power=False
        )
        rates_opt_opt_users[i] = calculate_system_rate(
            current_opt, h_hap2user, current_pairs,
            fixed_P, N_max, noise_linear, B, use_optimal_power=True
        )
    plt.figure(figsize=(10, 6))

    # 绘制三条曲线
    plt.plot(user_counts, rates_nearest_avg_users, 'b-o', linewidth=2, label='就近匹配 + 固定功率分配')
    plt.plot(user_counts, rates_opt_avg_users, 'r-s', linewidth=2, label='优化匹配 + 固定功率分配')
    plt.plot(user_counts, rates_opt_opt_users, 'g-d', linewidth=2, label='优化匹配 + 优化功率分配')

    # 添加标题和坐标轴标签
    plt.title('系统总和速率随用户对数量变化 (SNR=20dB)', fontsize=14)
    plt.xlabel('实用户对数量', fontsize=12)
    plt.ylabel('总和速率 (bps/Hz)', fontsize=12)

    # 显示图例和网格
    plt.legend(fontsize=12)
    plt.grid(True, linestyle='--', alpha=0.7)

    # 自动调整布局并显示图像
    plt.tight_layout()
    # 完整的输出文件路径
    output_filename = os.path.join(output_dir, "plot_users_SNR.png")
    # 保存图像
    plt.savefig(output_filename, dpi=300, bbox_inches='tight')
    plt.show()
    return [
        os.path.join(output_dir, "plot_users.png"),
        os.path.join(output_dir, "plot_HAP.png"),
        os.path.join(output_dir, "plot_users_SNR.png")
    ]


